{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Copyright (c) 2024, NVIDIA CORPORATION.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "#"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# Distributed Hyperparameter Tuning: Optuna + Spark Dataframes\n",
    "\n",
    "\n",
    "This demo demonstrates distributed hyperparameter tuning for XGBoost using Spark Dataframes.  \n",
    "We implement best practices to precompute data and maximize computations on the GPU.  \n",
    "\n",
    "Reference: https://forecastegy.com/posts/xgboost-hyperparameter-tuning-with-optuna/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Note:\n",
    "Before running, please make sure you've followed the relevant [setup instructions](../README.md) for your environment (standalone or databricks).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterable, List, Dict, Optional, Union, Sequence, Any\n",
    "import math\n",
    "import os\n",
    "import requests\n",
    "import pandas as pd\n",
    "import optuna\n",
    "from optuna.samplers import TPESampler\n",
    "import xgboost as xgb\n",
    "from pyspark.sql import SparkSession, DataFrame\n",
    "from pyspark import TaskContext, SparkConf\n",
    "from pyspark.sql.types import StructType, StructField, DoubleType, IntegerType, StringType, BooleanType"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download the dataset\n",
    "\n",
    "We'll use the [red wine quality dataset](https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv) to regress wine quality based on features such as acidity, sugar content, etc.  \n",
    "\n",
    "**Note**: This example uses a small dataset for demonstration purposes. The performance advantages of distributed training are best realized with large datasets and computational workloads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "cwd = os.getcwd()\n",
    "os.mkdir(os.path.join(cwd, \"data\")) if not os.path.exists(os.path.join(cwd, \"data\")) else None\n",
    "filepath = os.path.join(cwd, \"data\", \"winequality-red.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File downloaded and saved to /home/rishic/Code/myforks/spark-rapids-examples/examples/ML+DL-Examples/Optuna-Spark/optuna-examples/data/winequality-red.csv\n"
     ]
    }
   ],
   "source": [
    "url = \"https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv\"\n",
    "\n",
    "response = requests.get(url)\n",
    "if response.status_code == 200:\n",
    "    with open(filepath, \"wb\") as f:\n",
    "        f.write(response.content)\n",
    "    print(f\"File downloaded and saved to {filepath}\")\n",
    "else:\n",
    "    print(f\"Failed to download the file. Status code: {response.status_code}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1. Running Optuna locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudf\n",
    "from cuml.metrics.regression import mean_squared_error\n",
    "from cuml.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>fixed acidity</th>\n",
       "      <th>volatile acidity</th>\n",
       "      <th>citric acid</th>\n",
       "      <th>residual sugar</th>\n",
       "      <th>chlorides</th>\n",
       "      <th>free sulfur dioxide</th>\n",
       "      <th>total sulfur dioxide</th>\n",
       "      <th>density</th>\n",
       "      <th>pH</th>\n",
       "      <th>sulphates</th>\n",
       "      <th>alcohol</th>\n",
       "      <th>quality</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>7.4</td>\n",
       "      <td>0.70</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1.9</td>\n",
       "      <td>0.076</td>\n",
       "      <td>11.0</td>\n",
       "      <td>34.0</td>\n",
       "      <td>0.9978</td>\n",
       "      <td>3.51</td>\n",
       "      <td>0.56</td>\n",
       "      <td>9.4</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7.8</td>\n",
       "      <td>0.88</td>\n",
       "      <td>0.00</td>\n",
       "      <td>2.6</td>\n",
       "      <td>0.098</td>\n",
       "      <td>25.0</td>\n",
       "      <td>67.0</td>\n",
       "      <td>0.9968</td>\n",
       "      <td>3.20</td>\n",
       "      <td>0.68</td>\n",
       "      <td>9.8</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>7.8</td>\n",
       "      <td>0.76</td>\n",
       "      <td>0.04</td>\n",
       "      <td>2.3</td>\n",
       "      <td>0.092</td>\n",
       "      <td>15.0</td>\n",
       "      <td>54.0</td>\n",
       "      <td>0.9970</td>\n",
       "      <td>3.26</td>\n",
       "      <td>0.65</td>\n",
       "      <td>9.8</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>11.2</td>\n",
       "      <td>0.28</td>\n",
       "      <td>0.56</td>\n",
       "      <td>1.9</td>\n",
       "      <td>0.075</td>\n",
       "      <td>17.0</td>\n",
       "      <td>60.0</td>\n",
       "      <td>0.9980</td>\n",
       "      <td>3.16</td>\n",
       "      <td>0.58</td>\n",
       "      <td>9.8</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>7.4</td>\n",
       "      <td>0.70</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1.9</td>\n",
       "      <td>0.076</td>\n",
       "      <td>11.0</td>\n",
       "      <td>34.0</td>\n",
       "      <td>0.9978</td>\n",
       "      <td>3.51</td>\n",
       "      <td>0.56</td>\n",
       "      <td>9.4</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   fixed acidity  volatile acidity  citric acid  residual sugar  chlorides  \\\n",
       "0            7.4              0.70         0.00             1.9      0.076   \n",
       "1            7.8              0.88         0.00             2.6      0.098   \n",
       "2            7.8              0.76         0.04             2.3      0.092   \n",
       "3           11.2              0.28         0.56             1.9      0.075   \n",
       "4            7.4              0.70         0.00             1.9      0.076   \n",
       "\n",
       "   free sulfur dioxide  total sulfur dioxide  density    pH  sulphates  \\\n",
       "0                 11.0                  34.0   0.9978  3.51       0.56   \n",
       "1                 25.0                  67.0   0.9968  3.20       0.68   \n",
       "2                 15.0                  54.0   0.9970  3.26       0.65   \n",
       "3                 17.0                  60.0   0.9980  3.16       0.58   \n",
       "4                 11.0                  34.0   0.9978  3.51       0.56   \n",
       "\n",
       "   alcohol  quality  \n",
       "0      9.4        5  \n",
       "1      9.8        5  \n",
       "2      9.8        5  \n",
       "3      9.8        6  \n",
       "4      9.4        5  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = cudf.read_csv(filepath, delimiter=\";\")\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare the train/validation sets. Precompute the Quantile DMatrix, which is used by histogram-based tree methods to save memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = data.iloc[:, :-1].values\n",
    "y = data[\"quality\"].values\n",
    "X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "Xy_train_qdm = xgb.QuantileDMatrix(X_train, y_train)  # Precompute Quantile DMatrix to avoid repeated quantization every trial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Objective function\n",
    "\n",
    "We define the objective and a hyperparameter search space to optimize via the `trial.suggest_` methods.  \n",
    "\n",
    "In each trial, new hyperparameters will be suggested based on previous results. See [optuna.trial.Trial](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.trial.Trial.html) API for a full list of functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    params = {\n",
    "        \"objective\": \"reg:squarederror\",\n",
    "        \"verbosity\": 0,\n",
    "        \"learning_rate\": trial.suggest_float(\"learning_rate\", 1e-3, 0.1, log=True),\n",
    "        \"max_depth\": trial.suggest_int(\"max_depth\", 1, 10),\n",
    "        \"subsample\": trial.suggest_float(\"subsample\", 0.05, 1.0),\n",
    "        \"colsample_bytree\": trial.suggest_float(\"colsample_bytree\", 0.05, 1.0),\n",
    "        \"min_child_weight\": trial.suggest_int(\"min_child_weight\", 1, 20),\n",
    "        \"tree_method\": \"gpu_hist\",\n",
    "        \"device\": \"cuda\",\n",
    "    }\n",
    "\n",
    "    booster = xgb.train(params=params, dtrain=Xy_train_qdm, num_boost_round=trial.suggest_int(\"num_boost_round\", 100, 500))\n",
    "    predictions = booster.inplace_predict(X_val)\n",
    "    rmse = mean_squared_error(y_val, predictions, squared=False).get()\n",
    "    \n",
    "    return rmse   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the study and optimize. By default, the study results will be stored in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-12-11 23:47:48,356] A new study created in memory with name: optuna-xgboost-local\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-12-11 23:47:48,724] Trial 0 finished with value: 0.6377619522504244 and parameters: {'learning_rate': 0.005611516415334507, 'max_depth': 10, 'subsample': 0.7453942447208348, 'colsample_bytree': 0.6187255599871848, 'min_child_weight': 4, 'num_boost_round': 162}. Best is trial 0 with value: 0.6377619522504244.\n",
      "[I 2024-12-11 23:47:49,676] Trial 1 finished with value: 0.6703788974319568 and parameters: {'learning_rate': 0.0013066739238053278, 'max_depth': 9, 'subsample': 0.6210592611560484, 'colsample_bytree': 0.7226689489062432, 'min_child_weight': 1, 'num_boost_round': 488}. Best is trial 0 with value: 0.6377619522504244.\n",
      "[I 2024-12-11 23:47:49,819] Trial 2 finished with value: 0.6181751362616256 and parameters: {'learning_rate': 0.04622589001020832, 'max_depth': 3, 'subsample': 0.2227337188467456, 'colsample_bytree': 0.22423428436076215, 'min_child_weight': 7, 'num_boost_round': 310}. Best is trial 2 with value: 0.6181751362616256.\n",
      "[I 2024-12-11 23:47:49,942] Trial 3 finished with value: 0.6698576232920956 and parameters: {'learning_rate': 0.007309539835912915, 'max_depth': 3, 'subsample': 0.6312602499862605, 'colsample_bytree': 0.18251916761943976, 'min_child_weight': 6, 'num_boost_round': 246}. Best is trial 2 with value: 0.6181751362616256.\n",
      "[I 2024-12-11 23:47:50,060] Trial 4 finished with value: 0.6704590546150145 and parameters: {'learning_rate': 0.008168455894760165, 'max_depth': 8, 'subsample': 0.23969009305044175, 'colsample_bytree': 0.538522716492931, 'min_child_weight': 12, 'num_boost_round': 118}. Best is trial 2 with value: 0.6181751362616256.\n",
      "[I 2024-12-11 23:47:50,214] Trial 5 finished with value: 0.6088806682631155 and parameters: {'learning_rate': 0.016409286730647923, 'max_depth': 2, 'subsample': 0.11179901333601554, 'colsample_bytree': 0.9514412603906666, 'min_child_weight': 20, 'num_boost_round': 424}. Best is trial 5 with value: 0.6088806682631155.\n",
      "[I 2024-12-11 23:47:50,289] Trial 6 finished with value: 0.7103495949713845 and parameters: {'learning_rate': 0.0040665633135147945, 'max_depth': 1, 'subsample': 0.700021375186549, 'colsample_bytree': 0.4681448690526212, 'min_child_weight': 3, 'num_boost_round': 298}. Best is trial 5 with value: 0.6088806682631155.\n",
      "[I 2024-12-11 23:47:50,693] Trial 7 finished with value: 0.7255199474722185 and parameters: {'learning_rate': 0.001171593739230706, 'max_depth': 10, 'subsample': 0.29584098252001606, 'colsample_bytree': 0.6793961701362828, 'min_child_weight': 7, 'num_boost_round': 308}. Best is trial 5 with value: 0.6088806682631155.\n",
      "[I 2024-12-11 23:47:50,858] Trial 8 finished with value: 0.6060010014477214 and parameters: {'learning_rate': 0.0123999678368461, 'max_depth': 2, 'subsample': 0.9711053963763306, 'colsample_bytree': 0.7863761821930588, 'min_child_weight': 19, 'num_boost_round': 458}. Best is trial 8 with value: 0.6060010014477214.\n",
      "[I 2024-12-11 23:47:51,199] Trial 9 finished with value: 0.6292433375858283 and parameters: {'learning_rate': 0.015696396388661146, 'max_depth': 10, 'subsample': 0.13406787694932354, 'colsample_bytree': 0.23618371929818793, 'min_child_weight': 1, 'num_boost_round': 230}. Best is trial 8 with value: 0.6060010014477214.\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study(study_name=\"optuna-xgboost-local\", sampler=TPESampler(seed=42))\n",
    "study.optimize(objective, n_trials=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best RMSE:  0.6060010014477214\n",
      "Best hyperparameters:  {'learning_rate': 0.0123999678368461, 'max_depth': 2, 'subsample': 0.9711053963763306, 'colsample_bytree': 0.7863761821930588, 'min_child_weight': 19, 'num_boost_round': 458}\n"
     ]
    }
   ],
   "source": [
    "trial = study.best_trial\n",
    "print(\"Best RMSE: \", trial.value)\n",
    "print(\"Best hyperparameters: \", trial.params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2. Distributed Optuna on Spark "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PySpark\n",
    "\n",
    "For standalone users, we need to create the Spark session with the Spark-Rapids plugin. For Databricks users, the Spark session will be preconfigured and this cell can be skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plugin file already exists. Skipping download.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "24/12/11 23:47:51 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "24/12/11 23:47:51 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "24/12/11 23:47:52 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "24/12/11 23:47:52 WARN RapidsPluginUtils: RAPIDS Accelerator 25.02.1 using cudf 25.02.1, private revision bd4e99e18e20234ee0c54f95f4b0bfce18a6255e\n",
      "24/12/11 23:47:52 WARN RapidsPluginUtils: RAPIDS Accelerator is enabled, to disable GPU support set `spark.rapids.sql.enabled` to false.\n"
     ]
    }
   ],
   "source": [
    "def get_rapids_jar():\n",
    "    SPARK_RAPIDS_VERSION = \"25.04.0\"\n",
    "    rapids_jar = f\"rapids-4-spark_2.12-{SPARK_RAPIDS_VERSION}.jar\"\n",
    "    if not os.path.exists(rapids_jar):\n",
    "        print(\"Downloading Spark Rapids jar\")\n",
    "        url = f\"https://repo1.maven.org/maven2/com/nvidia/rapids-4-spark_2.12/{SPARK_RAPIDS_VERSION}/{rapids_jar}\"\n",
    "        response = requests.get(url)\n",
    "        if response.status_code == 200:\n",
    "            with open(rapids_jar, \"wb\") as f:\n",
    "                f.write(response.content)\n",
    "            print(f\"File '{rapids_jar}' downloaded and saved successfully.\")\n",
    "        else:\n",
    "            print(f\"Failed to download the plugin. Status code: {response.status_code}\")\n",
    "    else:\n",
    "        print(\"Plugin file already exists. Skipping download.\")\n",
    "    return rapids_jar\n",
    "\n",
    "def initialize_spark(rapids_jar: str):\n",
    "    import socket\n",
    "    hostname = socket.gethostname()\n",
    "    conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "\n",
    "    conf = SparkConf()\n",
    "    conf.setMaster(f\"spark://{hostname}:7077\")  # Assuming master is on host and default port. \n",
    "    conf.set(\"spark.task.maxFailures\", \"1\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", f\"{1/4}\")  # Setting to 1/4 for single-node demo. In practice, set to 1. \n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "    conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "    conf.set(\"spark.jars\", rapids_jar)\n",
    "    conf.set(\"spark.executorEnv.PYTHONPATH\", rapids_jar)\n",
    "    conf.set(\"spark.rapids.memory.gpu.minAllocFraction\", \"0.0001\")\n",
    "    conf.set(\"spark.plugins\", \"com.nvidia.spark.SQLPlugin\")\n",
    "    conf.set(\"spark.locality.wait\", \"0s\")\n",
    "    conf.set(\"spark.sql.cache.serializer\", \"com.nvidia.spark.ParquetCachedBatchSerializer\")\n",
    "    conf.set(\"spark.rapids.memory.gpu.pooling.enabled\", \"false\")\n",
    "    conf.set(\"spark.sql.execution.sortBeforeRepartition\", \"false\")\n",
    "    conf.set(\"spark.rapids.sql.format.parquet.reader.type\", \"MULTITHREADED\")\n",
    "    conf.set(\"spark.rapids.sql.format.parquet.multiThreadedRead.maxNumFilesParallel\", \"20\")\n",
    "    conf.set(\"spark.rapids.sql.multiThreadedRead.numThreads\", \"20\")\n",
    "    conf.set(\"spark.rapids.sql.python.gpu.enabled\", \"true\")\n",
    "    conf.set(\"spark.rapids.memory.pinnedPool.size\", \"2G\")\n",
    "    conf.set(\"spark.python.daemon.module\", \"rapids.daemon\")\n",
    "    conf.set(\"spark.rapids.sql.batchSizeBytes\", \"512m\")\n",
    "    conf.set(\"spark.sql.adaptive.enabled\", \"false\")\n",
    "    conf.set(\"spark.sql.files.maxPartitionBytes\", \"512m\")\n",
    "    conf.set(\"spark.rapids.sql.concurrentGpuTasks\", \"2\")\n",
    "    conf.set(\"spark.rapids.sql.explain\", \"NONE\")\n",
    "    \n",
    "    spark = SparkSession.builder.appName(\"optuna-spark-xgboost\").config(conf=conf).getOrCreate()\n",
    "    return spark\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    rapids_jar = get_rapids_jar()\n",
    "    spark = initialize_spark(rapids_jar)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Helper Class\n",
    "\n",
    "First we'll define a helper class. This will store the hyperparameters we want optimized in each trial, and easily convert that into a schema for the output dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OptunaParams:\n",
    "    def __init__(self):\n",
    "        self.hyperparameters = {}\n",
    "\n",
    "    def add_categorical_param(self, name: str, choices: Sequence[Union[None, bool, int, float, str]]):\n",
    "        \"\"\"\n",
    "        Adds a categorical hyperparameter to be tuned via Optuna's trial.suggest_categorical().\n",
    "        \"\"\"\n",
    "        self.hyperparameters[name] = { \"type\": \"categorical\", \"choices\": choices }\n",
    "    \n",
    "    def add_int_param(self, name: str, low: int, high: int, step: int = 1, log: bool = False):\n",
    "        \"\"\"\n",
    "        Adds an integer hyperparameter to be tuned via Optuna's trial.suggest_int().\n",
    "        \"\"\"\n",
    "        self.hyperparameters[name] = { \"type\": \"int\", \"low\": low, \"high\": high, \"step\": step, \"log\": log }\n",
    "    \n",
    "    def add_float_param(self, name: str, low: float, high: float, step: Optional[float] = None, log: bool = False):\n",
    "        \"\"\"\n",
    "        Adds a float hyperparameter to be tuned via Optuna's trial.suggest_float().\n",
    "        \"\"\"\n",
    "        self.hyperparameters[name] = { \"type\": \"float\", \"low\": low, \"high\": high, \"step\": step,\"log\": log }\n",
    "\n",
    "    def suggest_params(self, trial) -> Dict[str, Union[int, float, str, bool]]:\n",
    "        \"\"\"\n",
    "        Converts the hyperparameter space into a dictionary of suggested values in Optuna format,\n",
    "        to be called within the objective function.\n",
    "        \"\"\"\n",
    "        suggested_params = {}\n",
    "        for name, config in self.hyperparameters.items():\n",
    "            if config[\"type\"] == \"categorical\":\n",
    "                suggested_params[name] = trial.suggest_categorical(name, config[\"choices\"])\n",
    "            elif config[\"type\"] == \"int\":\n",
    "                suggested_params[name] = trial.suggest_int(\n",
    "                    name, config[\"low\"], config[\"high\"], step=config[\"step\"], log=config[\"log\"]\n",
    "                )\n",
    "            elif config[\"type\"] == \"float\":\n",
    "                suggested_params[name] = trial.suggest_float(\n",
    "                    name, config[\"low\"], config[\"high\"], step=config.get(\"step\", None), log=config[\"log\"]\n",
    "                )\n",
    "        return suggested_params\n",
    "\n",
    "    def to_schema(self) -> StructType:\n",
    "        \"\"\"\n",
    "        Converts the hyperparameter space into a Spark StructType output schema.\n",
    "        \"\"\"\n",
    "        fields = []\n",
    "        for name, config in self.hyperparameters.items():\n",
    "            if config[\"type\"] == \"float\":\n",
    "                fields.append(StructField(name, DoubleType(), False))\n",
    "            elif config[\"type\"] == \"int\":\n",
    "                fields.append(StructField(name, IntegerType(), False))\n",
    "            elif config[\"type\"] == \"categorical\":\n",
    "                if isinstance(config[\"choices\"][0], str):\n",
    "                    fields.append(StructField(name, StringType(), False))\n",
    "                elif isinstance(config[\"choices\"][0], bool):\n",
    "                    fields.append(StructField(name, BooleanType(), False))\n",
    "                elif isinstance(config[\"choices\"][0], (int, float)):\n",
    "                    fields.append(StructField(name, DoubleType(), False))\n",
    "                else:\n",
    "                    raise ValueError(f\"Unsupported categorical type for field {name}\")\n",
    "        \n",
    "        # Study will also return the best achieved loss:\n",
    "        fields.append(StructField(\"best_value\", DoubleType(), False)) \n",
    "        return StructType(fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optuna Task\n",
    "\n",
    "This implementation demonstrates **Spark I/O**.\n",
    "\n",
    "This means that Spark will read the dataset and create a duplicate of the dataset for each worker (1 partition = 1 duplicate), then map the tuning task onto each partition.  \n",
    "In practice, this enables the code to be chained to other Dataframe operations (e.g. ETL stages) without the intermediate step of writing to DBFS, at the cost of some overhead during duplication.\n",
    "\n",
    "For the alternative implementation using **Worker I/O**, see the [JoblibSpark notebook](optuna-joblibspark.ipynb). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the task, each worker will:\n",
    "1. Concatenate the pandas partition batches to form the dataset\n",
    "2. Load the study from the MySQL storage backend\n",
    "3. Optimize over the objective for the assigned number of trials, sending results back to the database after each iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def task_udf(pdf_iter: Iterable[pd.DataFrame],\n",
    "             xgb_params: Dict[str, Any],\n",
    "             optuna_params: OptunaParams,\n",
    "             trials_per_task: List[int],            \n",
    "             driver_ip: str,\n",
    "             study_name: str,\n",
    "             seed: int) -> Iterable[pd.DataFrame]:\n",
    "\n",
    "    import cudf\n",
    "    from cuml.metrics.regression import mean_squared_error\n",
    "    from cuml.model_selection import train_test_split\n",
    "    \n",
    "    tc = TaskContext.get()\n",
    "    assert \"gpu\" in tc.resources(), \"GPU resource not found.\"\n",
    "    num_trials = trials_per_task[tc.partitionId()]\n",
    "\n",
    "    df_list = []\n",
    "    for pdf in pdf_iter:\n",
    "        df_list.append(cudf.DataFrame.from_pandas(pdf))\n",
    "    \n",
    "    data = cudf.concat(df_list)\n",
    "    X = data.iloc[:, :-1].values\n",
    "    y = data[\"quality\"].values\n",
    "    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)\n",
    "\n",
    "    tuning_max_bin = \"max_bin\" in optuna_params.hyperparameters\n",
    "    if not tuning_max_bin:\n",
    "        max_bin = xgb_params.get(\"max_bin\", 256)\n",
    "        # Precompute Quantile DMatrix to avoid repeated quantization every trial.\n",
    "        Xy_train_qdm = xgb.QuantileDMatrix(X_train, y_train, max_bin=max_bin)\n",
    "\n",
    "    def objective(trial):\n",
    "        tuning_params = optuna_params.suggest_params(trial)\n",
    "        xgb_params.update(tuning_params)\n",
    "\n",
    "        if tuning_max_bin:\n",
    "            # If tuning the max_bin param, we must recompute the QDM every trial, since the quantiles change.\n",
    "            if \"n_estimators\" not in xgb_params:\n",
    "                xgb_params[\"n_estimators\"] = 100  # Default value if not tuning.\n",
    "\n",
    "            model = xgb.XGBRegressor(**xgb_params)\n",
    "            model.fit(X_train, y_train)\n",
    "            booster = model.get_booster()\n",
    "        else:\n",
    "            # Train the model with xgb.train() API using the precomputed QDM.\n",
    "            num_boost_round = xgb_params.get(\"n_estimators\", 100)\n",
    "            booster = xgb.train(params=xgb_params, dtrain=Xy_train_qdm, num_boost_round=num_boost_round)\n",
    "        \n",
    "        predictions = booster.inplace_predict(X_val)\n",
    "        rmse = mean_squared_error(y_val, predictions, squared=False).get()\n",
    "        \n",
    "        return rmse\n",
    "\n",
    "    study = optuna.load_study(\n",
    "        study_name=study_name,\n",
    "        storage=f\"mysql://optuna_user:optuna_password@{driver_ip}/optuna\",\n",
    "        sampler=TPESampler(seed=seed),\n",
    "    )\n",
    "\n",
    "    print(f\"Running {num_trials} trials on partition {tc.partitionId()}.\")\n",
    "    study.optimize(objective, n_trials=num_trials)\n",
    "\n",
    "    result_dict = {f\"{key}\": [value] for key, value in study.best_params.items()}\n",
    "    result_dict['best_value'] = [study.best_value]\n",
    "    \n",
    "    yield pd.DataFrame(result_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and run the Optuna study"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the driver IP for the MySQL database.  \n",
    "- For standalone users, make sure you've followed the [database setup instructions](../README.md#setup-database-for-optuna). The database should be on 'localhost'. \n",
    "- For databricks users, the database should already be setup on the driver node by the init script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if we're running on databricks\n",
    "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MySQL database is hosted on localhost\n"
     ]
    }
   ],
   "source": [
    "if on_databricks:\n",
    "    driver_ip = spark.conf.get(\"spark.driver.host\")\n",
    "else:\n",
    "    driver_ip = \"localhost\"\n",
    "\n",
    "print(f\"MySQL database is hosted on {driver_ip}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a new study, referencing the MySQL database as the storage backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-12-11 23:47:53,347] A new study created in RDB with name: optuna-xgboost-dataframe\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<optuna.study.study.Study at 0x756423c12560>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "study_name = \"optuna-xgboost-dataframe\"\n",
    "seed = 42\n",
    "\n",
    "try:\n",
    "    # Delete the study if it already exists\n",
    "    optuna.delete_study(\n",
    "        study_name=study_name, \n",
    "        storage=f\"mysql://optuna_user:optuna_password@{driver_ip}/optuna\"\n",
    "    )\n",
    "except:\n",
    "    pass\n",
    "\n",
    "optuna.create_study(\n",
    "    study_name=study_name,\n",
    "    storage=f\"mysql://optuna_user:optuna_password@{driver_ip}/optuna\",\n",
    "    sampler=TPESampler(seed=seed)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the number of tasks, number of trials, and trials per task. \n",
    "\n",
    "**NOTE**: for standalone users running on a single worker, the 4 tasks will all be assigned to the same worker and will time-share the GPU for demonstration. In practice, you should set `spark.task.resource.gpu.amount=1` and set num_tasks to the number of workers in the cluster so that each task gets full access to the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def partition_trials(total_trials: int, total_tasks: int) -> List[int]:\n",
    "    base_size = total_trials // total_tasks\n",
    "    extra = total_trials % total_tasks\n",
    "    partitions = [base_size] * total_tasks\n",
    "    for i in range(extra):\n",
    "        partitions[i] += 1\n",
    "    \n",
    "    return partitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trials per task: [25, 25, 25, 25]\n"
     ]
    }
   ],
   "source": [
    "num_tasks = 4\n",
    "num_trials = 100\n",
    "trials_per_task = partition_trials(num_trials, num_tasks)\n",
    "print(f\"Trials per task: {trials_per_task}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define params\n",
    "Define the XGBoost model params and the hyperparams for Optuna to tune. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Keep these params consistent:\n",
    "xgb_params = {\n",
    "    \"objective\": \"reg:squarederror\",\n",
    "    \"verbosity\": 0,\n",
    "    \"tree_method\": \"gpu_hist\",\n",
    "    \"device\": \"cuda\",\n",
    "    \"seed\": seed,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tune these params:\n",
    "hyperparams = OptunaParams()\n",
    "hyperparams.add_int_param(\"n_estimators\", low=100, high=500)\n",
    "hyperparams.add_float_param(\"learning_rate\", low=1e-3, high=0.1, log=True)\n",
    "hyperparams.add_int_param(\"max_depth\", low=1, high=10)\n",
    "hyperparams.add_float_param(\"subsample\", low=0.05, high=1.0)\n",
    "hyperparams.add_float_param(\"colsample_bytree\", low=0.05, high=1.0)\n",
    "hyperparams.add_int_param(\"min_child_weight\", low=1, high=20)\n",
    "\n",
    "out_schema = hyperparams.to_schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also define the following helper function, which will create duplicates of the dataframe held in separate partitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coalesce_tree_union(df: DataFrame, num_duplicates: int):\n",
    "    \"\"\"\n",
    "    Coalesce the DataFrame to a single partition and recursively self-union to create duplicates.\n",
    "    \"\"\"\n",
    "    input_df = df.coalesce(1).cache()\n",
    "    current_df = input_df\n",
    "    \n",
    "    if num_duplicates <= 1:\n",
    "        return current_df\n",
    "\n",
    "    recursions = int(math.log(num_duplicates, 2))\n",
    "    remainder = num_duplicates - 2 ** recursions\n",
    "\n",
    "    for _ in range(recursions):\n",
    "        current_df = current_df.union(current_df)\n",
    "\n",
    "    for _ in range(remainder):\n",
    "        current_df = current_df.union(input_df)\n",
    "    \n",
    "    return current_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load dataset\n",
    "\n",
    "Read the data from the local directory with Spark and then duplicate it to prepare to run the task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "if on_databricks:\n",
    "    # once the dataset is in dbfs, databricks appends \"dbfs:\" to the filepath automatically\n",
    "    filepath = '/FileStore/optuna-data/winequality-red.csv'\n",
    "else:\n",
    "    cwd = os.getcwd()\n",
    "    filepath = os.path.join(cwd, \"data\", \"winequality-red.csv\")\n",
    "\n",
    "in_schema = StructType([\n",
    "    StructField(\"fixed acidity\", DoubleType(), True),\n",
    "    StructField(\"volatile acidity\", DoubleType(), True),\n",
    "    StructField(\"citric acid\", DoubleType(), True),\n",
    "    StructField(\"residual sugar\", DoubleType(), True),\n",
    "    StructField(\"chlorides\", DoubleType(), True),\n",
    "    StructField(\"free sulfur dioxide\", DoubleType(), True),\n",
    "    StructField(\"total sulfur dioxide\", DoubleType(), True),\n",
    "    StructField(\"density\", DoubleType(), True),\n",
    "    StructField(\"pH\", DoubleType(), True),\n",
    "    StructField(\"sulphates\", DoubleType(), True),\n",
    "    StructField(\"alcohol\", DoubleType(), True),\n",
    "    StructField(\"quality\", IntegerType(), True)\n",
    "])\n",
    "\n",
    "data_df = spark.read.csv(filepath, header=True, schema=in_schema, sep=\";\")\n",
    "data_df = coalesce_tree_union(data_df, num_duplicates=num_tasks)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the study\n",
    "\n",
    "Map the Optuna task onto the dataframe and collect the results (it might take a few minutes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "result_df = data_df.mapInPandas(lambda pdf_iter: \n",
    "                                task_udf(pdf_iter,\n",
    "                                         xgb_params=xgb_params,\n",
    "                                         optuna_params=hyperparams,\n",
    "                                         trials_per_task=trials_per_task,\n",
    "                                         driver_ip=driver_ip,\n",
    "                                         study_name=study_name,\n",
    "                                         seed=seed),\n",
    "                                         schema=out_schema).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best parameters: {'n_estimators': 419.0, 'learning_rate': 0.015039610889407229, 'max_depth': 10.0, 'subsample': 0.6630214978050138, 'colsample_bytree': 0.8524338650689898, 'min_child_weight': 2.0}\n",
      "Best value: 0.533100375625104\n"
     ]
    }
   ],
   "source": [
    "results = result_df.iloc[0].to_dict()\n",
    "best_value = results.pop(\"best_value\")\n",
    "\n",
    "print(f\"Best parameters: {results}\")\n",
    "print(f\"Best value: {best_value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "optuna-spark",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
